package edu.washington.ling.syntric.proc;

import org.apache.commons.lang.StringUtils;
import org.kohsuke.args4j.CmdLineException;
import org.kohsuke.args4j.CmdLineParser;

import edu.washington.ling.syntric.io.*;
import edu.washington.ling.syntric.struct.error.InvalidValueException;
import edu.washington.ling.syntric.struct.error.PrintingExceptionHandler;
import edu.washington.ling.syntric.struct.error.UnreadableRepresentationException;
import edu.washington.ling.syntric.struct.tree.Token;
import edu.washington.ling.syntric.struct.tree.Tree;

import java.io.IOException;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.LinkedList;
import java.util.Map;


public class CreateUnifiedExamples {
	private static CreateUnifiedExamplesOptions opts = new CreateUnifiedExamplesOptions();
	private static PrintingExceptionHandler pHandler = new PrintingExceptionHandler();
	
	private static String getYieldString(Tree tree) {
		List<String> words = new LinkedList<String>();
		for (Token tok: tree.getTokenList()) {
			if (!tok.isEmptyCategory())
				words.add(tok.getWord());
		}
		return StringUtils.join(words, " ");
	}

	public static void addExistingExamples(ExampleDoc exf, 
			Map<String, Example> exrMap) throws IOException {
		Example exr = null;
		while ((exr = exf.getNextExample()) != null) {
			if (!exrMap.containsKey(exr.getID())) {
				exrMap.put(exr.getID(), exr);
			} 
		}
	}
	public static void addTrees(TreeFileReader tfr, TreeReader trr, 
			Map<String, Example> exrMap) throws IOException {
		List<String> ids = new LinkedList<String>(tfr.getIDs());
		String format = trr.getFormat();
		String layer = tfr.getLayer();
		for (String id: ids) {
			TreeRepresentation tr = tfr.readTreeRepWithID(id);
			if (tr == null) {
				continue;
			}
			tr.setFormat(format);
			if (exrMap.containsKey(id)) {
				Example exr = exrMap.get(id);
				try {
					Tree tree = trr.readTree(tr);
					String yield = getYieldString(tree);
					if (exr.hasOriginal()) {
						if(!exr.getOriginal().trim().equals(yield.trim())) {
							tr.addMetaData("tokenDiff");
							String errStr = String.format("Error: Could not add the '%s' tree for sentence '%s'.\r\n" +
									"The stored sentence does not match the current tree yield.", layer, id);
							System.err.println(errStr);
							System.err.println(String.format("%20s: %s", "Stored sentence", exr.getOriginal()));
							System.err.println(String.format("%20s: %s", "Current tree yield", yield));
							System.err.println();
						}
					}
					else {
						exr.setOriginal(yield);
					}
				}
				catch(UnreadableRepresentationException e) {
					tr.addMetaData("unreadable");
					pHandler.handle(e, String.format("Tree '%s' is unreadable", id));
				}
				finally {
					exr.addTreeRep(tr);
				}
			}
			else {
				Example exr = new Example(id);
				try {
					Tree tree = trr.readTree(tr);
					String yield = getYieldString(tree);
					exr.setOriginal(yield);
				}
				catch(UnreadableRepresentationException e) {
					tr.addMetaData("unreadable");
					pHandler.handle(e, String.format("Tree '%s' is unreadable", id));
				}
				finally {
					exr.addTreeRep(tr);
					exrMap.put(id, exr);
				}
			}
		}
	}
	
	private static Map<String, Example> createUnifiedExamples() {
		Map<String, Example> exrMap = new HashMap<String, Example>();
		try {
			if (opts.inputFilePath != null && opts.inputFileFormat != null) {
				ExampleDocReader exfr = ExampleDocReader.newReader(opts.inputFilePath, opts.inputFileFormat);
				ExampleDoc exf = exfr.read();
				addExistingExamples(exf, exrMap);
			}
			if (opts.dsFilePath != null && opts.dsFileFormat != null && opts.dsTreeFormat != null) {
				TreeFileReader tfr = TreeFileReader.newReader(opts.dsFilePath, opts.dsFileFormat, "DS");
				TreeReader trr = TreeReader.newReader(opts.dsTreeFormat);
				addTrees(tfr, trr, exrMap);
			}
			if (opts.pbFilePath != null && opts.pbFileFormat != null && opts.pbTreeFormat != null) {
				TreeFileReader tfr = TreeFileReader.newReader(opts.pbFilePath, opts.pbFileFormat, "PB");
				TreeReader trr = TreeReader.newReader(opts.pbTreeFormat);
				addTrees(tfr, trr, exrMap);
			}
			if (opts.psFilePath != null && opts.psFileFormat != null && opts.psTreeFormat != null) {
				TreeFileReader tfr = TreeFileReader.newReader(opts.psFilePath, opts.psFileFormat, "PS");
				TreeReader trr = TreeReader.newReader(opts.psTreeFormat);
				addTrees(tfr, trr, exrMap);
			}
		}
		catch(InvalidValueException e) {
			pHandler.handle(e, "Failed to create unified examples", System.out);
			System.exit(-1);
		}
		catch (IOException e) {
			pHandler.handle(e, "Failed to create unified examples", System.out);
			e.printStackTrace(System.out);
			System.exit(-1);
		} 
		return exrMap;
	}
	
	private static void printExampleRepsToFile(Collection<Example> exrs) {
		int numExamples = 0;
		try {
			ExampleDocWriter exfw = ExampleDocWriter.newWriter(opts.outputFilePath, opts.outputFileFormat);
			for (Example exr: exrs) {
				numExamples++;
				exfw.writeExample(exr);
			}
			exfw.close();
			System.out.println("Print unified examples was successful.");
			System.out.println(String.format("Processed %d total examples", numExamples));
		}
		catch (IOException e) {
			pHandler.handle(e, "Failed to print unified examples", System.out);
			e.printStackTrace(System.out);
			System.exit(-1);
		}
	}
	public static void main(String[] args) {
		CmdLineParser parser = new CmdLineParser(opts);
		try {
			parser.parseArgument(args);
			Map<String, Example> exrMap = createUnifiedExamples();
			printExampleRepsToFile(exrMap.values());
		}
		catch(CmdLineException e) {
			System.out.println(e.getMessage());
			System.out.print("CreateUnifiedExamples");
			parser.printSingleLineUsage(System.out);
			System.out.println();
			parser.printUsage(System.out);
			System.exit(-1);
		}
	}
}
